#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
自然语言搜索工具
功能：支持自然语言问题搜索，自动提取关键词并进行智能匹配
"""

import re
import jieba
import jieba.posseg as pseg
from pymilvus import connections, Collection
from sklearn.feature_extraction.text import TfidfVectorizer
import sys


class NaturalLanguageSearch:
    def __init__(self, collection_name="knowledge_lines", host="localhost", port="19530"):
        self.collection_name = collection_name
        self.host = host
        self.port = port
        self.collection = None
        self.vectorizer = None
        self.is_connected = False
        
        # 基本停用词（最小集合）
        self.stop_words = {'的', '是', '在', '有', '和', '怎么', '什么', '吗', '呢', '了'}

        # 产品编号模式
        self.product_pattern = re.compile(r'[A-Z]{1,3}\d{3,4}')
    
    def connect_to_milvus(self):
        """连接到Milvus数据库"""
        try:
            connections.connect("default", host=self.host, port=self.port)
            self.collection = Collection(self.collection_name)
            self.collection.load()
            print(f"✅ 连接成功: {self.collection_name} (共{self.collection.num_entities}条记录)")
            self.is_connected = True
            return True
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return False
    
    def extract_keywords(self, text):
        """从自然语言中提取关键词 - 简化版，只做分词"""
        # 1. 提取产品编号
        products = self.product_pattern.findall(text.upper())

        # 2. 简单分词
        try:
            # 使用jieba分词
            words = list(jieba.cut(text))
            keywords = []

            for word in words:
                # 过滤停用词、单字符和纯数字
                if (word not in self.stop_words and
                    len(word) > 1 and
                    not word.isdigit() and
                    word.strip()):
                    keywords.append(word)

        except Exception as e:
            print(f"⚠️  分词失败: {e}")
            keywords = []

        # 3. 合并产品编号和关键词
        all_keywords = products + keywords

        # 4. 去重并保持顺序
        unique_keywords = []
        seen = set()
        for kw in all_keywords:
            if kw not in seen:
                unique_keywords.append(kw)
                seen.add(kw)

        return unique_keywords, products
    
    def build_search_expressions(self, keywords, products):
        """构建搜索表达式"""
        expressions = []
        
        # 1. 如果有产品编号，优先搜索该产品
        if products:
            for product in products:
                # 精确匹配产品编号
                expressions.append(f'line_text like "%{product}%"')
        
        # 2. 关键词搜索
        for keyword in keywords:
            if keyword not in products and len(keyword) > 1:
                expressions.append(f'line_text like "%{keyword}%"')
        
        return expressions
    
    def search_with_expressions(self, expressions, limit=10):
        """使用表达式搜索 - 简化版"""
        if not expressions:
            return []

        try:
            # 使用OR连接所有表达式
            expr = ' or '.join(expressions)

            results = self.collection.query(
                expr=expr,
                limit=limit * 2,  # 获取更多结果用于排序
                output_fields=['filename', 'line_number', 'line_text']
            )

            return results

        except Exception as e:
            print(f"❌ 搜索失败: {e}")
            return []
    
    def rank_results(self, results, keywords, products):
        """对搜索结果进行排序 - 简化版"""
        scored_results = []

        for result in results:
            text = result['line_text'].lower()
            score = 0

            # 产品编号匹配得分
            for product in products:
                if product.lower() in text:
                    score += 10

            # 关键词匹配得分
            for keyword in keywords:
                if keyword.lower() in text:
                    score += 5
                    # 完整词匹配额外加分
                    if f' {keyword.lower()} ' in f' {text} ':
                        score += 2

            # 文本长度惩罚（更短的文本可能更相关）
            score -= len(text) / 1000

            scored_results.append((score, result))

        # 按分数排序
        scored_results.sort(key=lambda x: x[0], reverse=True)

        return [result for score, result in scored_results]
    
    def natural_search(self, query, limit=5):
        """自然语言搜索主函数"""
        print(f"🤖 自然语言搜索: '{query}'")
        print("=" * 60)
        
        # 1. 提取关键词
        keywords, products = self.extract_keywords(query)
        print(f"🔍 提取的关键词: {keywords}")
        if products:
            print(f"📦 识别的产品: {products}")
        
        # 2. 构建搜索表达式
        expressions = self.build_search_expressions(keywords, products)
        if not expressions:
            print("❌ 无法提取有效的搜索关键词")
            return []
        
        print(f"🔧 搜索策略: {len(expressions)} 个条件")
        
        # 3. 执行搜索
        results = self.search_with_expressions(expressions, limit * 2)  # 获取更多结果用于排序
        
        if not results:
            print("❌ 没有找到匹配的结果")
            return []
        
        # 4. 结果排序
        ranked_results = self.rank_results(results, keywords, products)
        
        # 5. 返回前N个结果
        final_results = ranked_results[:limit]
        print(f"✅ 找到 {len(final_results)} 条相关结果")
        
        return final_results
    
    def display_results(self, results, query):
        """显示搜索结果"""
        if not results:
            print("\n❌ 没有找到相关结果")
            return
        
        print(f"\n🎯 搜索结果:")
        print("=" * 60)
        
        for i, result in enumerate(results, 1):
            print(f"\n--- 结果 {i} ---")
            print(f"📁 文件: {result['filename']}")
            print(f"📍 行号: {result['line_number']}")
            
            # 高亮显示关键词
            text = result['line_text']
            try:
                keywords, products = self.extract_keywords(query)
                # 确保keywords和products都是列表
                if not isinstance(keywords, list):
                    keywords = []
                if not isinstance(products, list):
                    products = []

                all_keywords = keywords + products
                for keyword in all_keywords:
                    if keyword and len(str(keyword)) > 1:
                        # 不区分大小写的替换
                        pattern = re.compile(re.escape(str(keyword)), re.IGNORECASE)
                        text = pattern.sub(f"【{keyword}】", text)
            except Exception as e:
                print(f"⚠️  关键词高亮失败: {e}")
                # 如果高亮失败，直接显示原文本
            
            print(f"📝 内容: {text}")
            
            if i >= 10:  # 最多显示10条
                print(f"\n... 还有 {len(results) - 10} 条结果未显示")
                break


def interactive_search():
    """交互式自然语言搜索"""
    print("🤖 自然语言智能搜索工具")
    print("=" * 50)
    print("💡 支持自然语言问题，如：")
    print("   - TD504的扩增循环数是怎么样的？")
    print("   - PCR扩增失败怎么办？")
    print("   - 如何提高转化效率？")
    print("=" * 50)
    
    # 初始化搜索器
    searcher = NaturalLanguageSearch()
    
    # 连接数据库
    if not searcher.connect_to_milvus():
        return
    
    print("✅ 系统准备就绪！")
    print("💡 提示：输入 'quit' 或 'exit' 退出程序")
    
    # 开始交互式搜索
    while True:
        print("\n" + "=" * 50)
        query = input("🤖 请输入您的问题: ").strip()
        
        if query.lower() in ['quit', 'exit', '退出']:
            print("👋 再见！")
            break
        
        if not query:
            print("⚠️  请输入有效的问题")
            continue
        
        try:
            # 执行自然语言搜索
            results = searcher.natural_search(query, limit=10)
            
            # 显示结果
            searcher.display_results(results, query)
            
        except KeyboardInterrupt:
            print("\n\n👋 用户中断，再见！")
            break
        except Exception as e:
            print(f"\n❌ 搜索过程中出现错误: {e}")


def command_line_search():
    """命令行搜索"""
    if len(sys.argv) < 2:
        print("用法: python natural_language_search.py '您的自然语言问题'")
        print("示例: python natural_language_search.py 'TD504的扩增循环数是怎么样的？'")
        return
    
    query = ' '.join(sys.argv[1:])
    
    # 初始化搜索器
    searcher = NaturalLanguageSearch()
    
    # 连接和搜索
    if not searcher.connect_to_milvus():
        return
    
    # 执行搜索
    results = searcher.natural_search(query, limit=10)
    searcher.display_results(results, query)


if __name__ == "__main__":
    if len(sys.argv) > 1:
        command_line_search()
    else:
        interactive_search()
